import numpy as np

class CustomKMeans:
    def __init__(self, n_clusters, max_iter=300, tol=1e-4, random_state=None):
        """
        自定义 K-Means 算法
        参数:
            n_clusters (int): 聚类数
            max_iter (int): 最大迭代次数
            tol (float): 收敛阈值
            random_state (int, optional): 随机种子
        """
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.tol = tol
        self.random_state = random_state

    def fit(self, X):
        """
        训练 K-Means 模型
        参数:
            X (ndarray): 数据集 (n_samples, n_features)
        """
        np.random.seed(self.random_state)
        n_samples, n_features = X.shape

        # 随机初始化质心
        initial_indices = np.random.choice(n_samples, self.n_clusters, replace=False)
        self.cluster_centers_ = X[initial_indices]

        for iteration in range(self.max_iter):
            # 计算每个点到每个质心的距离
            distances = np.linalg.norm(X[:, np.newaxis] - self.cluster_centers_, axis=2)

            # 分配每个点到最近的质心
            self.labels_ = np.argmin(distances, axis=1)

            # 计算新的质心
            new_centers = np.array([X[self.labels_ == j].mean(axis=0) for j in range(self.n_clusters)])

            # 检查是否收敛
            if np.all(np.linalg.norm(new_centers - self.cluster_centers_, axis=1) < self.tol):
                break

            self.cluster_centers_ = new_centers

        # 计算惯性值（样本点到最近质心的距离平方和）
        self.inertia_ = np.sum((X - self.cluster_centers_[self.labels_])**2)

    def predict(self, X):
        """
        对新数据点进行预测
        参数:
            X (ndarray): 数据集 (n_samples, n_features)
        返回:
            labels (ndarray): 每个点的聚类标签
        """
        distances = np.linalg.norm(X[:, np.newaxis] - self.cluster_centers_, axis=2)
        return np.argmin(distances, axis=1)
